// SPDX-License-Identifier: Mulan PSL v2
/*
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This software is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *         http://license.coscl.org.cn/MulanPSL2
 *
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

use crate::api::*;

use indexmap::{indexmap, IndexMap};
use lazy_static::lazy_static;
use thiserror::Error;

use xgpu_common::{ipc::message::Argument, utils::api_name::ApiFuncName};

#[derive(Debug, Error)]
pub enum ServerErr {
    #[error("API run failed, api: {0}, error code: {1}")]
    ApiRunError(String, i32),

    #[error("Invalid argument type: {0}")]
    InvalidType(String),
}

pub trait ApiHandler: Send + Sync {
    fn handle_api(&self, args: &mut [Argument<'_>]) -> Result<Argument<'static>, ServerErr>;
}

lazy_static! {
    pub static ref FUNC_HANDLER_MAP: IndexMap<u64, Box<dyn ApiHandler>> = indexmap! {
        (ApiFuncName::FuncCudadevicesynchronize as u64) => Box::new(CudaDeviceSynchronizeHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCudadevicegetstreampriorityrange as u64) => Box::new(CudaDeviceGetStreamPriorityRangeHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCudagetlasterror as u64) => Box::new(CudaGetLastErrorHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCudapeekatlasterror as u64) => Box::new(CudaPeekAtLastErrorHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCudagetdevicecount as u64) => Box::new(CudaGetDeviceCountHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCudagetdevicepropertiesV2 as u64) => Box::new(CudaGetDevicePropertiesV2Handler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCudadevicegetattribute as u64) => Box::new(CudaDeviceGetAttributeHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCudasetdevice as u64) => Box::new(CudaSetDeviceHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCudagetdevice as u64) => Box::new(CudaGetDeviceHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCudastreamcreatewithpriority as u64) => Box::new(CudaStreamCreateWithPriorityHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCudathreadexchangestreamcapturemode as u64) => Box::new(CudaThreadExchangeStreamCaptureModeHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCudastreamiscapturing as u64) => Box::new(CudaStreamIsCapturingHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCudastreamwaitevent as u64) => Box::new(CudaStreamWaitEventHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCudastreamsynchronize as u64) => Box::new(CudaStreamSynchronizeHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncNcclgetversion as u64) => Box::new(NcclGetVersionHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCudaeventcreatewithflags as u64) => Box::new(CudaEventCreateWithFlagsHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCudaeventrecord as u64) => Box::new(CudaEventRecordHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCudaeventrecordwithflags as u64) => Box::new(CudaEventRecordWithFlagsHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCudaeventquery as u64) => Box::new(CudaEventQueryHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCudaeventsynchronize as u64) => Box::new(CudaEventSynchronizeHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCudaeventdestroy as u64) => Box::new(CudaEventDestroyHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCudaeventelapsedtime as u64) => Box::new(CudaEventElapsedTimeHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCudamemset as u64) => Box::new(CudaMemsetHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCudapointergetattributes as u64) => Box::new(CudaPointerGetAttributesHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCudeviceget as u64) => Box::new(CuDeviceGetHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCudevicegetattribute as u64) => Box::new(CuDeviceGetAttributeHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCudeviceprimaryctxgetstate as u64) => Box::new(CuDevicePrimaryCtxGetStateHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCuctxgetcurrent as u64) => Box::new(CuCtxGetCurrentHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCufuncgetattribute as u64) => Box::new(CuFuncGetAttributeHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCudamemsetasync as u64) => Box::new(CudaMemsetAsyncHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCudafuncsetattribute as u64) => Box::new(CudaFuncSetAttributeHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCudafuncgetattributes as u64) => Box::new(CudaFuncGetAttributesHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCublascreateV2 as u64) => Box::new(CublasCreateV2handler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCublassetworkspaceV2 as u64) => Box::new(CublasSetWorkspaceV2handler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCublassetstreamV2 as u64) => Box::new(CublasSetStreamV2handler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCublasgetstreamV2 as u64) => Box::new(CublasGetStreamV2handler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCublasltmatmulpreferencecreate as u64) => Box::new(CublasLtMatmulPreferenceCreateHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCublasltmatmuldesccreate as u64) => Box::new(CublasLtMatmulDescCreateHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCublasltmatmuldescsetattribute as u64) => Box::new(CublasLtMatmulDescSetAttributeHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCublasltmatrixlayoutcreate as u64) => Box::new(CublasLtMatrixLayoutCreateHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCublasltmatmulpreferencesetattribute as u64) => Box::new(CublasLtMatmulPreferenceSetAttributeHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCublasgetmathmode as u64) => Box::new(CublasGetMathModeHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCublassetmathmode as u64) => Box::new(CublasSetMathModeHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCublasltmatmulalgogetheuristic as u64) => Box::new(CublasLtMatmulAlgoGetHeuristicHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCublasltmatmul as u64) => Box::new(CublasLtMatmulHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCublasltmatrixlayoutdestroy as u64) => Box::new(CublasLtMatrixLayoutDestroyHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCublasltmatmuldescdestroy as u64) => Box::new(CublasLtMatmulDescDestroyHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCublasltmatmulpreferencedestroy as u64) => Box::new(CublasLtMatmulPreferenceDestroyHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCublassgemmV2 as u64) => Box::new(CublasSgemmV2handler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCudamemcpyasync as u64) => Box::new(CudaMemcpyAsyncHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCudamalloc as u64) => Box::new(CudaMallocHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCudafree as u64) => Box::new(CudaFreeHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncNvmlinitV2 as u64) => Box::new(NvmlInitV2Handler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncNvmlshutdown as u64) => Box::new(NvmlShutdownHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncNvmldevicegetcountV2 as u64) => Box::new(NvmlDeviceGetCountV2handler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncNcclcommdestroy as u64) => Box::new(NcclCommDestroyHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncNcclcommabort as u64) => Box::new(NcclCommAbortHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncNcclcommfinalize as u64) => Box::new(NcclCommFinalizeHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncNcclcommsplit as u64) => Box::new(NcclCommSplitHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncNcclcommgetasyncerror as u64) => Box::new(NcclCommGetAsyncErrorHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncNcclbcast as u64) => Box::new(NcclBcastHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncNcclallreduce as u64) => Box::new(NcclAllReduceHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncNcclallgather as u64) => Box::new(NcclAllGatherHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncNcclsend as u64) => Box::new(NcclSendHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncNcclrecv as u64) => Box::new(NcclRecvHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncNcclgroupstart as u64) => Box::new(NcclGroupStartHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncNcclgroupend as u64) => Box::new(NcclGroupEndHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncNcclgetuniqueid as u64) => Box::new(NcclGetUniqueIdHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncNcclcomminitrankconfig as u64) => Box::new(NcclCommInitRankConfigHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCudalaunchkernel as u64) => Box::new(CudaLaunchKernelHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncNcclreduce as u64) => Box::new(NcclReduceHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCublasgemmex as u64) => Box::new(CublasGemmExHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCublassgemmstridedbatched as u64) => Box::new(CublasSgemmStridedBatchedHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCublasgemmstridedbatchedex as u64) => Box::new(CublasGemmStridedBatchedExHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCublasltcreate as u64) => Box::new(CublasLtCreateHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCublasltdestroy as u64) => Box::new(CublasLtDestroyHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncCudaoccupancymaxactiveblockspermultiprocessorwithflags as u64) => Box::new(CudaOccupancyMaxActiveBlocksPerMultiprocessorWithFlagsHandler) as Box<dyn ApiHandler>,
        (ApiFuncName::FuncNcclcomminitrankconfigrecovery as u64) => Box::new(NcclCommInitRankConfigRecoveryHandler) as Box<dyn ApiHandler>,

    };
}

pub fn call_handler(key: u64, args: &mut [Argument<'_>]) -> Result<Argument<'static>, ServerErr> {
    match FUNC_HANDLER_MAP.get(&key) {
        Some(handler) => handler.handle_api(args),
        None => Err(ServerErr::ApiRunError(
            "call_handler.get failed, api_enum:".to_string(),
            key as i32,
        )),
    }
}
